"""
将图片替换到PDF脚本（覆盖原PDF）
根据mask文件夹中的文件名，在panel文件夹中找到对应图片，然后根据json文件内对应坐标信息替换到PDF中。
生成的PDF将直接覆盖各文件夹内的同名原PDF。
"""

import os
import io
import json
import glob
from PIL import Image
import fitz  # PyMuPDF

# ========= 配置 =========
SPLICING_RESULTS_DIR = "D:\\FILES\\科研实践\\ICLR 2026\\results_copy-move"  # 根目录
UPDATED_FORGERY_JSON = "D:\\FILES\\科研实践\\ICLR 2026\\results_copy-move\\updated_forgery_copy_move.json"  # JSON文件路径
OUTPUT_DIR = "D:\\FILES\科研实践\\ICLR 2026\\results_copy-move"  # 输出目录

# 处理模式选择
BATCH_MODE = True  # 设置为True进行批量处理，False处理单个文件夹
BATCH_SKIP_EXISTING = True  # 批量处理时是否跳过已存在的输出文件

# 单文件夹处理模式配置（仅在BATCH_MODE=False时使用）
# 注意：批量处理模式下此配置无效，脚本会自动扫描所有文件夹
SELECTED_FOLDER = "2010072670_FG_001"  # 例如: "2010072785_FG_001"

PAGE_PX_PER_PT = 1.0  # 像素到点的转换比率

# ========= 基础工具 =========
def safe_overwrite_pdf(doc: fitz.Document, out_pdf: str) -> None:
    """安全覆盖保存PDF，避免repaired PDF报错"""
    try:
        if hasattr(doc, "can_save_incrementally") and doc.can_save_incrementally():
            doc.save(out_pdf, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
            return
        doc.save(out_pdf, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
        return
    except Exception:
        tmp = out_pdf + ".tmp"
        doc.save(tmp, garbage=4, deflate=True, encryption=fitz.PDF_ENCRYPT_KEEP)
        try:
            os.replace(tmp, out_pdf)
        except Exception:
            import shutil
            shutil.move(tmp, out_pdf)

def crop_and_resize_to_bytes(ori_img_path: str, out_w: int, out_h: int) -> bytes:
    """将图片裁剪并调整到指定尺寸，返回字节流"""
    with Image.open(ori_img_path) as img:
        # 保持宽高比的情况下居中裁剪
        target_aspect = out_w / float(out_h) if out_h != 0 else 1
        w, h = img.size
        
        if w / float(h) > target_aspect:
            new_w = int(round(h * target_aspect))
            x0 = max(0, (w - new_w) // 2)
            img = img.crop((x0, 0, x0 + new_w, h))
        else:
            new_h = int(round(w / target_aspect)) if target_aspect != 0 else h
            y0 = max(0, (h - new_h) // 2)
            img = img.crop((0, y0, w, y0 + new_h))
        
        # 调整到精确尺寸
        img = img.convert("RGB").resize((max(1, out_w), max(1, out_h)), Image.Resampling.LANCZOS)
        
        # 转换为字节流
        bio = io.BytesIO()
        img.save(bio, format="PNG", optimize=True)
        return bio.getvalue()

def page_pixel_to_pt(bbox_px: list) -> list:
    """将像素坐标转换为点坐标"""
    k = PAGE_PX_PER_PT
    return [bbox_px[0] / k, bbox_px[1] / k, bbox_px[2] / k, bbox_px[3] / k]

def load_json_data(json_path: str) -> dict:
    """加载JSON数据文件"""
    try:
        with open(json_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as e:
        print(f"错误：无法加载JSON文件 {json_path}: {e}")
        return {}

def find_panel_info_in_json(json_data: dict, panel_filename: str):
    """
    根据panel文件名在JSON中查找对应的坐标信息
    
    参数:
        json_data: JSON数据
        panel_filename: panel文件名（如 "2010072785_FG_001_7_1_1.png"）
    
    返回:
        (bbox, page_index, paper_id): 边界框、页面索引和论文ID
    """
    # 从文件名提取panel_id（去掉扩展名）
    panel_id = os.path.splitext(panel_filename)[0]
    
    papers = json_data.get("papers", [])
    for paper in papers:
        paper_id = paper.get("paper_id", "")
        pages = paper.get("pages", [])
        
        for page in pages:
            page_index = page.get("page_index", 1)
            figures = page.get("figures", [])
            
            for fig in figures:
                panels = fig.get("panels", [])
                for panel in panels:
                    if panel.get("panel_id") == panel_id:
                        bbox = panel.get("bbox_page_pixel", [])
                        if len(bbox) == 4:
                            print(f"找到panel信息: {panel_id}")
                            print(f"  论文ID: {paper_id}")
                            print(f"  页面: {page_index}")
                            print(f"  坐标: {bbox}")
                            return bbox, page_index - 1, paper_id  # 转换为0-based页面索引
    
    print(f"错误：在JSON中未找到panel {panel_id}")
    return None, None, None

def get_mask_files_with_subdirs(mask_dir: str) -> list:
    """获取mask目录下的所有图片文件，返回(子文件夹名, 文件名)的列表"""
    mask_files = []
    if os.path.exists(mask_dir):
        # 遍历mask目录下的子文件夹
        for subdir in os.listdir(mask_dir):
            subdir_path = os.path.join(mask_dir, subdir)
            if os.path.isdir(subdir_path):
                # 遍历子文件夹中的文件
                for file in os.listdir(subdir_path):
                    if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                        mask_files.append((subdir, file))
    return mask_files

def find_corresponding_panel_file(panel_dir: str, mask_subdir: str, mask_filename: str) -> str:
    """在panel目录中根据子文件夹名和文件名查找对应的文件"""
    # 首先在panel目录中查找与mask子文件夹名相同的文件夹
    panel_subdir_path = os.path.join(panel_dir, mask_subdir)
    if os.path.exists(panel_subdir_path):
        # 在该子文件夹中查找与mask文件名相同的文件
        panel_file_path = os.path.join(panel_subdir_path, mask_filename)
        if os.path.exists(panel_file_path):
            return panel_file_path
    
    print(f"警告：在panel目录中未找到 {mask_subdir}/{mask_filename}")
    return None

def replace_single_image_in_doc(doc: fitz.Document, image_path: str, bbox_px: list, page_index: int):
    """
    在已打开的PDF文档中替换单张图片
    
    参数:
        doc: 已打开的PDF文档对象
        image_path: 要插入的图片路径
        bbox_px: 替换区域的边界框 [x0, y0, x1, y1] (像素单位)
        page_index: 页面索引(0-based)
    """
    # 检查页面索引是否有效
    if page_index < 0 or page_index >= len(doc):
        raise ValueError(f"页面索引 {page_index} 无效，文档只有 {len(doc)} 页")
    
    # 获取指定页面
    page = doc[page_index]
    
    # 将像素坐标转换为点坐标
    bbox_pt = page_pixel_to_pt(bbox_px)
    rect = fitz.Rect(bbox_pt)
    
    # 计算替换区域的尺寸
    w_overlay = int(round(bbox_px[2] - bbox_px[0]))
    h_overlay = int(round(bbox_px[3] - bbox_px[1]))
    
    if w_overlay <= 0 or h_overlay <= 0:
        raise ValueError("边界框尺寸无效")
    
    # 准备要插入的图片
    overlay_bytes = crop_and_resize_to_bytes(image_path, w_overlay, h_overlay)
    
    # 插入图片到指定区域
    page.insert_image(rect, stream=overlay_bytes, overlay=True, keep_proportion=False)
    
    print(f"成功替换图片到第{page_index + 1}页: {os.path.basename(image_path)}")

def replace_multiple_images_in_pdf(pdf_path: str, replacement_info_list: list, output_path: str):
    """
    在PDF中替换多张图片，生成一个包含所有修改的PDF
    
    参数:
        pdf_path: 原始PDF文件路径
        replacement_info_list: 替换信息列表，每个元素包含 (image_path, bbox_px, page_index)
        output_path: 输出PDF路径
    """
    # 打开PDF文档
    doc = fitz.open(pdf_path)
    
    try:
        # 处理每个替换信息
        for i, (image_path, bbox_px, page_index) in enumerate(replacement_info_list):
            print(f"处理第 {i+1}/{len(replacement_info_list)} 张图片...")
            replace_single_image_in_doc(doc, image_path, bbox_px, page_index)
        
        # 保存包含所有修改的文档
        safe_overwrite_pdf(doc, output_path)
        print(f"成功生成包含 {len(replacement_info_list)} 张图片替换的PDF: {output_path}")
        
    finally:
        doc.close()

def get_all_folders_to_process():
    """获取splicing/results下所有需要处理的文件夹"""
    folders = []
    if os.path.exists(SPLICING_RESULTS_DIR):
        for item in os.listdir(SPLICING_RESULTS_DIR):
            item_path = os.path.join(SPLICING_RESULTS_DIR, item)
            if os.path.isdir(item_path) and item != "output":  # 排除output目录
                # 检查是否包含必要的子目录和文件
                mask_dir = os.path.join(item_path, "mask")
                panel_dir = os.path.join(item_path, "panel")
                pdf_file = os.path.join(item_path, f"{item}.pdf")
                
                if (os.path.exists(mask_dir) and 
                    os.path.exists(panel_dir) and 
                    os.path.exists(pdf_file)):
                    folders.append(item)
    return sorted(folders)

def process_single_folder(folder_name, json_data):
    """处理单个文件夹"""
    print(f"\n{'='*60}")
    print(f"开始处理文件夹: {folder_name}")
    print(f"{'='*60}")
    
    # 构建路径
    selected_folder_path = os.path.join(SPLICING_RESULTS_DIR, folder_name)
    mask_dir = os.path.join(selected_folder_path, "mask")
    panel_dir = os.path.join(selected_folder_path, "panel")
    pdf_path = os.path.join(selected_folder_path, f"{folder_name}.pdf")
    
    # 检查必要目录和文件是否存在
    if not os.path.exists(selected_folder_path):
        print(f"错误：选择的文件夹不存在: {selected_folder_path}")
        return False
    
    if not os.path.exists(mask_dir):
        print(f"错误：mask文件夹不存在: {mask_dir}")
        return False
    
    if not os.path.exists(panel_dir):
        print(f"错误：panel文件夹不存在: {panel_dir}")
        return False
    
    if not os.path.exists(pdf_path):
        print(f"错误：PDF文件不存在: {pdf_path}")
        return False
    
    # 直接覆盖原PDF
    output_pdf_path = pdf_path
    
    print(f"PDF路径: {pdf_path}")
    print(f"Mask目录: {mask_dir}")
    print(f"Panel目录: {panel_dir}")
    
    # 获取mask文件列表（包含子文件夹信息）
    mask_files = get_mask_files_with_subdirs(mask_dir)
    if not mask_files:
        print("错误：mask目录中没有找到图片文件")
        return False
    
    print(f"找到 {len(mask_files)} 个mask文件:")
    for subdir, filename in mask_files:
        print(f"  {subdir}/{filename}")
    
    # 收集所有需要替换的图片信息
    replacement_info_list = []
    success_count = 0
    
    for mask_subdir, mask_file in mask_files:
        print(f"\n处理mask文件: {mask_subdir}/{mask_file}")
        
        # 在panel目录中查找对应的文件
        panel_file_path = find_corresponding_panel_file(panel_dir, mask_subdir, mask_file)
        if not panel_file_path:
            print(f"警告：在panel目录中未找到对应的文件: {mask_subdir}/{mask_file}")
            continue
        
        print(f"找到对应的panel文件: {panel_file_path}")
        
        # 在JSON中查找坐标信息
        bbox_px, page_index, paper_id = find_panel_info_in_json(json_data, mask_file)
        if not bbox_px:
            print(f"警告：在JSON中未找到 {mask_file} 的坐标信息")
            continue
        
        # 添加到替换信息列表
        replacement_info_list.append((panel_file_path, bbox_px, page_index))
        success_count += 1
        print(f"已收集替换信息: 第{page_index + 1}页, 坐标{bbox_px}")
    
    # 如果有需要替换的图片，执行批量替换
    if replacement_info_list:
        print(f"\n开始批量替换 {len(replacement_info_list)} 张图片...")
        
        try:
            # 执行批量图片替换
            replace_multiple_images_in_pdf(
                pdf_path=pdf_path,
                replacement_info_list=replacement_info_list,
                output_path=output_pdf_path
            )
            print(f"处理完成！成功替换 {success_count}/{len(mask_files)} 个文件到单个PDF中")
            print(f"已覆盖PDF: {output_pdf_path}")
            return True
        except Exception as e:
            print(f"错误：批量替换时发生异常: {e}")
            return False
    else:
        print("没有找到任何可以替换的图片")
        return False

def main():
    # 检查JSON文件是否存在
    if not os.path.exists(UPDATED_FORGERY_JSON):
        print(f"错误：JSON文件不存在: {UPDATED_FORGERY_JSON}")
        return
    
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    # 加载JSON数据
    json_data = load_json_data(UPDATED_FORGERY_JSON)
    if not json_data:
        return
    
    if BATCH_MODE:
        # 批量处理模式
        print("="*80)
        print("批量处理模式")
        print("="*80)
        
        # 获取所有需要处理的文件夹
        folders_to_process = get_all_folders_to_process()
        if not folders_to_process:
            print("错误：没有找到任何可处理的文件夹")
            return
        
        print(f"找到 {len(folders_to_process)} 个可处理的文件夹:")
        for i, folder in enumerate(folders_to_process, 1):
            print(f"  {i}. {folder}")
        
        # 统计信息
        total_folders = len(folders_to_process)
        success_folders = 0
        skipped_folders = 0
        failed_folders = 0
        
        # 处理每个文件夹
        for i, folder_name in enumerate(folders_to_process, 1):
            print(f"\n进度: {i}/{total_folders}")
            try:
                result = process_single_folder(folder_name, json_data)
                if result:
                    success_folders += 1
                else:
                    failed_folders += 1
            except Exception as e:
                print(f"处理文件夹 {folder_name} 时发生异常: {e}")
                failed_folders += 1
        
        # 输出最终统计
        print("\n" + "="*80)
        print("批量处理完成！")
        print("="*80)
        print(f"总文件夹数: {total_folders}")
        print(f"成功处理: {success_folders}")
        print(f"处理失败: {failed_folders}")
        print("所有PDF已覆盖至各自原路径。")
        
    else:
        # 单文件夹处理模式
        print("="*80)
        print("单文件夹处理模式")
        print("="*80)
        
        result = process_single_folder(SELECTED_FOLDER, json_data)
        if result:
            print("\n处理完成！PDF已覆盖原路径。")
        else:
            print(f"\n处理失败！")

if __name__ == "__main__":
    main()